package com.cl.ks.utils;

import org.springframework.web.multipart.MultipartFile;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;

import com.fasterxml.jackson.databind.ObjectMapper;

public class FileUploadSecurity {

    // 允许的文件类型白名单
    private static final Set<String> ALLOWED_EXTENSIONS = new HashSet<>(Arrays.asList("jpg", "jpeg", "png", "gif", "json"));
    private static final long MAX_FILE_SIZE = 5 * 1024 * 1024; // 5MB

    // 文件头信息（魔数）映射
    private static final Map<String, byte[]> FILE_SIGNATURES = new HashMap<>();

    static {
        {
            FILE_SIGNATURES.put("jpg", new byte[]{(byte) 0xFF, (byte) 0xD8, (byte) 0xFF});
            FILE_SIGNATURES.put("jpeg", new byte[]{(byte) 0xFF, (byte) 0xD8, (byte) 0xFF});
            FILE_SIGNATURES.put("png", new byte[]{(byte) 0x89, (byte) 0x50, (byte) 0x4E, (byte) 0x47, (byte) 0x0D, (byte) 0x0A, (byte) 0x1A, (byte) 0x0A});
            FILE_SIGNATURES.put("gif", new byte[]{(byte) 'G', (byte) 'I', (byte) 'F', (byte) '8'});
        }
    }

    /**
     * 检查文件扩展名是否在允许的列表内。
     */
    public static boolean isValidExtension(String fileName) {
        String extension = getFileExtension(fileName).toLowerCase();
        return ALLOWED_EXTENSIONS.contains(extension);
    }

    /**
     * 获取文件扩展名。
     */
    private static String getFileExtension(String fileName) {
        int lastIndexOfDot = fileName.lastIndexOf('.');
        if (lastIndexOfDot == -1 || lastIndexOfDot == 0 || lastIndexOfDot == fileName.length() - 1) {
            return "";
        }
        return fileName.substring(lastIndexOfDot + 1);
    }

    /**
     * 检查文件大小是否超过最大限制。
     */
    public static boolean isWithinSizeLimit(File file) {
        return file.length() <= MAX_FILE_SIZE;
    }

    /**
     * 验证文件内容类型（MIME类型）。
     */
    public static boolean isValidContentType(File file) throws IOException {
        Path path = Paths.get(file.getAbsolutePath());
        String contentType = Files.probeContentType(path);
        if (contentType == null) {
            return false;
        }
        String fileExtension = getFileExtension(file.getName()).toLowerCase();
        switch (fileExtension) {
            case "jpg":
            case "jpeg":
                return contentType.startsWith("image/jpeg");
            case "png":
                return contentType.startsWith("image/png");
            case "gif":
                return contentType.startsWith("image/gif");
            case "json":
                return contentType.equals("application/json");
            default:
                return false;
        }
    }

    /**
     * 验证文件头信息（魔数）。对于JSON文件不做此检查。
     */
    public static boolean isValidFileSignature(File file) throws IOException {
        String fileExtension = getFileExtension(file.getName()).toLowerCase();
        if ("json".equals(fileExtension)) {
            return true; // 不需要检查JSON文件的魔数
        }

        byte[] headerBytes = new byte[8]; // 假设最大的文件头长度为8个字节
        try (InputStream inputStream = Files.newInputStream(file.toPath())) {
            int bytesRead = inputStream.read(headerBytes);
            if (bytesRead < 1) {
                return false;
            }
        }

        byte[] expectedSignature = FILE_SIGNATURES.get(fileExtension);
        if (expectedSignature == null) {
            return false;
        }

        // 只检查文件头中实际存在的字节数
        for (int i = 0; i < Math.min(expectedSignature.length, headerBytes.length); i++) {
            if (headerBytes[i] != expectedSignature[i]) {
                return false;
            }
        }
        return true;
    }

    /**
     * 校验JSON文件的有效性。
     */
    public static boolean isValidJsonFile(File file) throws IOException {
        ObjectMapper objectMapper = new ObjectMapper();
        try {
            objectMapper.readTree(file);
            return true;
        } catch (IOException e) {
            return false;
        }
    }

    /**
     * 生成安全的随机文件名。
     */
    public static String generateSecureFileName(String originalFilename) {
        String extension = getFileExtension(originalFilename);
        String randomName = java.util.UUID.randomUUID().toString();
        return randomName + (extension.isEmpty() ? "" : "." + extension);
    }

    public static void validateFile(MultipartFile multipartFile) throws IOException {
        if (multipartFile.isEmpty()) {
            throw new IllegalArgumentException("The file is empty.");
        }

        // 将MultipartFile转换为临时文件以便后续操作
        File tempFile = convertMultipartFileToFile(multipartFile);

        try {
            // 执行安全检查
            String fileExtension = getFileExtension(multipartFile.getOriginalFilename()).toLowerCase();
            if (!isValidExtension(multipartFile.getOriginalFilename()) ||
                    !isWithinSizeLimit(tempFile) ||
                    !isValidContentType(tempFile) ||
                    (!"json".equals(fileExtension) && !isValidFileSignature(tempFile)) ||
                    ("json".equals(fileExtension) && !isValidJsonFile(tempFile))) {
                throw new SecurityException("Invalid file upload attempt detected.");
            }
        } finally {
            // 清理临时文件
            if (tempFile.exists()) {
                tempFile.delete();
            }
        }
    }
    /**
     * 统一验证并保存MultipartFile。
     */
    public static void validateAndSaveMultipartFile(MultipartFile multipartFile, String uploadDir) throws IOException {
        if (multipartFile.isEmpty()) {
            throw new IllegalArgumentException("The file is empty.");
        }

        // 将MultipartFile转换为临时文件以便后续操作
        File tempFile = convertMultipartFileToFile(multipartFile);

        try {
            // 执行安全检查
            String fileExtension = getFileExtension(multipartFile.getOriginalFilename()).toLowerCase();
            if (!isValidExtension(multipartFile.getOriginalFilename()) ||
                    !isWithinSizeLimit(tempFile) ||
                    !isValidContentType(tempFile) ||
                    (!"json".equals(fileExtension) && !isValidFileSignature(tempFile)) ||
                    ("json".equals(fileExtension) && !isValidJsonFile(tempFile))) {
                throw new SecurityException("Invalid file upload attempt detected.");
            }

            // 生成安全文件名并保存文件
            String secureFileName = generateSecureFileName(multipartFile.getOriginalFilename());
            Path targetLocation = Paths.get(uploadDir, secureFileName);
            Files.move(tempFile.toPath(), targetLocation);

        } finally {
            // 清理临时文件
            if (tempFile.exists()) {
                tempFile.delete();
            }
        }
    }

    /**
     * 将MultipartFile转换为临时文件。
     */
    private static File convertMultipartFileToFile(MultipartFile multipartFile) throws IOException {
        File tempFile = Files.createTempFile(null, null).toFile();
        multipartFile.transferTo(tempFile);
        return tempFile;
    }
}
